# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# import numpy as np


# DETR Transformer class.
# Copy-paste from torch.nn.Transformer with modifications:
#     * positional encodings are passed in MHattention
#     * extra LN at the end of encoder is removed
#     * decoder returns a stack of activations from all decoding layers
# """
import copy
from typing import Optional, List

import torch
import torchvision
from torch.nn import functional as F
from torch import nn, Tensor


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")


class TransformerEncoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(src, pos)
        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        return src

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)


class MyModel(nn.Module):
    def __init__(self, num_classes):
        super().__init__()

        resnet18 = torchvision.models.resnet18(pretrained=True)
        children = list(resnet18.children())

        self.pre_conv = nn.Sequential(*children[0:4])
        
        self.block_1 = children[4]
        self.block_2 = children[5]
        self.block_3 = children[6]
        self.block_4 = children[7]

        self.att1 = TransformerEncoderLayer(d_model=56*56, nhead=4, dim_feedforward=128)
        self.att2 = TransformerEncoderLayer(d_model=28*28, nhead=4, dim_feedforward=128)
        self.att3 = TransformerEncoderLayer(d_model=14*14, nhead=4, dim_feedforward=128)
        self.att4 = TransformerEncoderLayer(d_model=7*7, nhead=7, dim_feedforward=128)

        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

    def forward(self, inputs):
        bs = inputs.size(0)

        outputs = self.pre_conv(inputs)

        outputs_block_1 = self.block_1(outputs)
        outputs_att_1 = self.att1(outputs_block_1.view(bs, outputs_block_1.size(1), -1))

#         print(outputs_att_1.size())
        outputs_1 = outputs_block_1 * outputs_att_1.view(bs, outputs_block_1.size(1), outputs_block_1.size(2), outputs_block_1.size(3))
        # print(outputs_att_1.size())

        outputs_block_2 = self.block_2(outputs_1)
        outputs_att_2 = self.att2(outputs_block_2.view(bs, outputs_block_2.size(1), -1))
        outputs_2 = outputs_block_2 * outputs_att_2.view(bs, outputs_block_2.size(1), outputs_block_2.size(2), outputs_block_2.size(3))
        # print(outputs_2.size())

        outputs_block_3 = self.block_3(outputs_2)
        outputs_att_3 = self.att3(outputs_block_3.view(bs, outputs_block_3.size(1), -1))
        outputs_3 = outputs_block_3 * outputs_att_3.view(bs, outputs_block_3.size(1), outputs_block_3.size(2), outputs_block_3.size(3))
        # print(outputs_3.size())

        outputs_block_4 = self.block_4(outputs_3)
        # outputs_att_4 = self.att4(outputs_block_4).view(bs, outputs_block_4.size(1), -1)
        # outputs_4 = outputs_block_4 * outputs_att_4.view(bs, outputs_block_4.size(1), outputs_block_4.size(2), outputs_block_4.size(3))
        outputs_4 = outputs_block_4

        outputs = self.gap(outputs_4).view(inputs.size(0), -1)
        outputs = self.fc(outputs)

        return outputs
        

# net = MyModel(num_classes=7)
# x = torch.randn((2, 3, 224, 224))
# out = net(x)
# print(out.size())
